//
//  human_detector.cpp
//  yuzon
//
//  Created by Alvin on 2023/3/29.
//

#include "human_detector.hpp"
#include "utils.hpp"

using namespace cv;
using namespace std;
using namespace paddle::lite_api;

HumanDetector::HumanDetector(const string &model_detection_path,
                             const string &label_list_path) :
input_width_(kHumanInputWidth),
input_height_(kHumanInputHeight),
threshold_(kHumanThreshold) {
  MobileConfig config;
  config.set_model_from_file(model_detection_path);
  config.set_power_mode(LITE_POWER_NO_BIND);
  config.set_threads(4);
  predictor_ = CreatePaddlePredictor<MobileConfig>(config);
  labels_ = LabelList(label_list_path);
  colors_ = ColorMap(labels_.size());
}

void HumanDetector::Predict(const Mat &image, HumanResults &results) {
  Preprocess(image);
  predictor_->Run();
  Postprocess(image, results);
}

void HumanDetector::Preprocess(const Mat &image) {
  const vector<string> &input_names = predictor_->GetInputNames();
  
  for (const string &input_name : input_names) {
    if (input_name == "image") {
      auto input_tensor = predictor_->GetInputByName(input_name);
      input_tensor->Resize({1, 3, input_height_, input_width_});
      auto input_data = input_tensor->mutable_data<float>();
      
      Mat input_image;
      resize(image, input_image, Size(input_width_, input_height_));
      input_image.convertTo(input_image, CV_32FC3);
      Permute(input_image, input_data);
    } else if (input_name == "scale_factor") {
      auto scale_tensor = predictor_->GetInputByName(input_name);
      scale_tensor->Resize({1, 2});
      auto scale_data = scale_tensor->mutable_data<float>();
      
      scale_data[0] = float(input_height_) / float(image.rows);
      scale_data[1] = float(input_width_) / float(image.cols);
    }
  }
}

void HumanDetector::Postprocess(const Mat &image, HumanResults &results) {
  auto output_tensor = predictor_->GetOutput(0);
  auto output_data = output_tensor->data<float>();
  
  auto bbox_tensor = predictor_->GetOutput(1);
  auto bbox_data = bbox_tensor->data<int>();
  
  for (int i = 0; i < bbox_data[0]; ++i) {
    int class_id = round(output_data[i * 6]);
    float score = output_data[1 + i * 6];
    
    if (score < threshold_)
      continue;
    
    int xmin = round(output_data[2 + i * 6]);
    int ymin = round(output_data[3 + i * 6]);
    int xmax = round(output_data[4 + i * 6]);
    int ymax = round(output_data[5 + i * 6]);
    
    int w = xmax - xmin;
    int h = ymax - ymin;
    
    Rect rect = Rect(xmin, ymin, w, h) & Rect(0, 0, image.cols, image.rows);
    
    if (rect.area() > 0) {
      HumanResult result{class_id, score, rect, labels_[class_id], colors_[class_id]};
      results.emplace_back(result);
    }
  }
  
  // 结果排序
  if (results.size() > 1)
    sort(results.begin(), results.end(), CompareHumanResult);
}
